import os
import sys
import time
from multiprocessing import Process

def SubmitJob(jobName, device_ids):
    log_line = f"[{time.strftime('%Y-%m-%d %H:%M:%S')}]: {jobName} start!"
    print(log_line)
    jobLog = jobName.replace(".py", ".log")
    if os.path.exists(jobLog):
        os.system("rm {}".format(jobLog))
    os.system(f"CUDA_VISIBLE_DEVICES={device_ids} python {jobName} >> {jobLog}")
    log_line = f"[{time.strftime('%Y-%m-%d %H:%M:%S')}]: {jobName} end!"
    print(log_line)

def collect_logs(job_list, log_suffix):
    log_dirs = [os.path.join(dir, fname.strip(".py")) for fname in job_list]
    final_log_dir = "./{}_Logs".format(log_suffix)
    if os.path.exists(final_log_dir):
        os.system(f"rm -r {final_log_dir}")
    os.system(f"mkdir {final_log_dir}")
    for dir_name in log_dirs:
        for fname in os.listdir(dir_name):
            file = os.path.join(dir_name, fname)
            if os.path.isdir(file):
                os.system(f"cp -r {file} {final_log_dir}/")
    os.system(f"zip -r {final_log_dir}.zip {final_log_dir}")
    os.system(f"rm -r {final_log_dir}")

dir = "./"
dir_path = os.path.abspath(dir)
files = os.listdir(dir_path)
py_files = [fname for fname in files if ("MetaSelfTrainV4_" in fname) and ("FS20.py" in fname)]
py_files.sort()

process_list = []
start_idx = 0
for device_id, idx in enumerate(range(start_idx, min(len(py_files), start_idx+8))):
    p = Process(target=SubmitJob, args=(py_files[idx], f'{device_id}', ))
    p.start()
    process_list.append(p)

for p in process_list:
    p.join()

collect_logs(job_list=py_files, log_suffix="MetaSelfTrainV4_FS20")